Add target-penalty annealing support (penalty_schedule) to the training loops#43
Merged
andrewrosemberg merged 7 commits intoJun 15, 2026
Conversation
…tiple_shooting (LearningToOptimize#41) A new penalty_schedule keyword accepts nothing (default, unchanged behavior), :default_annealed, or an explicit vector of (first_batch, last_batch, multiplier) phases. Multipliers scale the objective coefficients of the norm_deficit variables relative to the penalty the models were built with, covering the L1-only, L2-only, and combined modes of create_deficit! as well as deterministic-equivalent and shooting-window model copies.
…y schedule (LearningToOptimize#41) The examples now make a single training call spanning num_epochs*num_batches batches so the schedule covers the whole run. This also keeps one optimizer state throughout (previously Flux.setup re-ran every epoch), makes a true return from the record callback stop the whole run, and drops the iter-2100 batch-doubling hook that the collapsed iteration count would have activated. Constant penalties remain available via penalty_schedule = nothing.
…r) into penalty-annealing branch # Conflicts: # examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl # src/DecisionRules.jl # src/simulate_multistage.jl # test/runtests.jl
Welcome to Codecov 🎉Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests. ℹ️ You can also turn on project coverage checks and project coverage reporting on Pull Request comment Thanks for integrating Codecov - We've got you covered ☂️ |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #41
Summary
Training with a constant target-slack penalty weight either hides target violations (low/moderate weight:
the deterministic-equivalent view looks great while the stage-wise rollout violates its own targets) or
degrades the learned policy (high weight from epoch 1). This PR adds the penalty-schedule option proposed
in #41: a multiplier ladder ×0.1 → ×1 → ×10 → ×30 applied during training, relative to the penalty the
models were built with (so with
penalty=:auto, relative to the auto-calibrated base= max |objective coefficient|).
What changed
src/utils.jl:default_annealed_schedule(num_batches)(phase lengths proportional to 2/2/4/16, everyphase ≥ 1 batch, < 4 batches falls back to the tail of the ladder so runs always end at ×30), schedule
validation (contiguous from batch 1, finite positive multipliers, past-end hold), and internal helpers
that capture the objective coefficients of the
norm_deficitvariables once at training start andrescale them at phase boundaries. Scaling that single coefficient covers all three
create_deficit!modes (L1-only, L2-only, and the combined mode where the L1/L2 weights live in the linking constraint),
as well as the renamed copies inside deterministic-equivalent and shooting-window models — no model
rebuild is needed.
src/simulate_multistage.jl/src/multiple_shooting.jl: a newpenalty_schedulekeyword on bothtrain_multistageoverloads and ontrain_multiple_shooting.nothing(default) leaves behaviorexactly as today;
:default_annealedbuilds the ladder overnum_batches; an explicit vector of(first_batch, last_batch, multiplier)phases is validated and used as given.examples/HydroPowerModels: the three trainers default to the annealed schedule. They now make a singletraining call spanning
num_epochs*num_batchesbatches so the schedule covers the whole run. Deliberateexample-behavior notes: this keeps one optimizer state for the whole run (previously
Flux.setupre-ranevery epoch, silently resetting Adam moments), makes a
truereturn from the record callback stop thewhole run (previously it only ended the current epoch), and drops the
iter % 2100batch-doubling hook(which the collapsed iteration count would have newly activated). Constant-penalty training remains
available via
penalty_schedule = nothing.test/runtests.jl: schedule-resolution exactness (including the 24-batch reference ladder1:2 ×0.1, 3:4 ×1, 5:8 ×10, 9:24 ×30), validation rejections, past-end hold, coefficient scaling in allthree
create_deficit!modes, untouched coefficients on the default path, and an end-to-endtrain_multistagerun crossing a phase boundary on DiffOpt models.Validation
Pkg.test()green on this branch (210/210, includes the new testsets).fixed 6-scenario held-out set): with the constant as-built penalty the trained policy ends at a
stage-wise-rollout target-violation share of 0.220 — far outside the ≤ ~0.05 trust gate — while its
deterministic-equivalent view looks unremarkable; with the annealed default at the same budget the share
drops to 0.052 — near but still slightly above the ≤ ~0.05 gate (4.2× lower). This is partial,
reduced-scale evidence for the core mechanism of the issue (constant penalties learn targets that are
not followable stage by stage; annealing restores followability); it is not reduced-scale gap closure.
The full gap-closure endpoint of the issue table (det-eq ≡ rollout to ~1e-7 relative) requires
full-scale training and is documented in the issue itself.
Notes
penalty_schedule=nothingis the default and the gated codepath leaves every objective coefficient untouched. Only the examples change their default to annealed.
by construction); annealing them separately is out of scope.
default_annealed_scheduleis exported; the capture/apply helpers are internal.examples/HydroPowerModels/train_dr_hydropowermodels_subproblems.jl; whichever merges second will berebased onto the other.